module grt_coefficients
  use parameters
  use matrix_operations
  implicit none
  
contains
  
  ! 计算反射/透射系数
  subroutine grt_coefs(kn, o)
    real(kind=8), intent(in) :: kn
    complex(kind=8), intent(in) :: o
    complex(kind=8) :: a(4,4), b(4,4), ab(4,4), ed(4,4), eu(4,4), cu, cd
    complex(kind=8) :: tem1, tem2, tem3
    real(kind=8) :: hu, hd
    integer :: lay, i, j
    
    ! 计算修改的R/T系数(j=1,2,...,N)
    ! SH情况
    do lay = 1, nly-1
      
      cu = csn(lay) * media%mu(lay)
      cd = csn(lay+1) * media%mu(lay+1)
      rt_coef%tdd0(lay) = 2.0d0 * cu / (cu + cd)
      rt_coef%rdu0(lay) = (cu - cd) / (cu + cd)
      
      rt_coef%rud0(lay) = -rt_coef%rdu0(lay)
      rt_coef%tuu0(lay) = 2.0d0 * cd / (cu + cd)
    end do
    
    
    
    ! PSV情况
    do lay = 1, nly-1
      
      call mtxe(lay, kn, o, eu)
      call mtxe(lay+1, kn, o, ed)

      ! print *, 001, lay,nly-1 ,o, media%vp(lay+1)
      ! print *, 001, lay,nly-1 ,o, csn(lay+1)
      ! print *, 001, lay,nly-1 ,o, cpn(lay+1)
      do j = 1, 2
        do i = 1, 2
          a(i,j) = ed(i,j)
          a(2+i,j) = ed(2+i,j)
          a(i,j+2) = -eu(i,j+2)
          a(i+2,j+2) = -eu(i+2,j+2)
          b(i,j) = eu(i,j)
          b(2+i,j) = eu(2+i,j)
          b(i,j+2) = -ed(i,j+2)
          b(i+2,j+2) = -ed(i+2,j+2)
        end do
      end do
      call mtxinv_4(a)
      
      call mtxmtp(4, a, b, ab)
      
      do j = 1, 2
        do i = 1, 2
          rt_coef%tdd(lay,i,j) = ab(i,j)
          rt_coef%rdu(lay,i,j) = ab(i+2,j)
          rt_coef%rud(lay,i,j) = ab(i,j+2)
          rt_coef%tuu(lay,i,j) = ab(i+2,j+2)
        end do
      end do
    end do
    
    ! 计算广义R/T系数
    ! 自由表面(lay=0)的GRUD
    ! SH波
    rt_coef%grud0(0) = (1.0d0, 0.0d0)
  
    ! PSV波
    call mtxe(1, kn, o, eu)
    do j = 1, 2
      do i = 1, 2
        a(i,j) = eu(i+2,j)
        b(i,j) = eu(i+2,j+2)
      end do
    end do
    call mtxinv_2(a)
    do j = 1, 2
      do i = 1, 2
        rt_coef%grud(0,i,j) = -(a(i,1)*b(1,j) + a(i,2)*b(2,j))
      end do
    end do
    
    ! lay=1,2,...,(ls-1)界面的GRUD和GTUU
    do lay = 1, ls-1
      hu = media%z(lay) - media%z(lay-1)
      exu(1) = exp(-cpn(lay)*hu)
      exu(2) = exp(-csn(lay)*hu)
      rt_coef%ex(lay,1) = exu(1)
      rt_coef%ex(lay,2) = exu(2)
      
      ! SH波
      tem1 = rt_coef%grud0(lay-1) * exu(2)
      tem2 = rt_coef%rdu0(lay) * exu(2)
      tem3 = rt_coef%tdd0(lay) * exu(2)
      rt_coef%gtu0(lay) = rt_coef%tuu0(lay) / (1.0d0 - tem1*tem2)
      rt_coef%grud0(lay) = rt_coef%rud0(lay) + tem3*tem1*rt_coef%gtu0(lay)
      
      ! PSV波
      do j = 1, 2
        do i = 1, 2
          a(i,j) = -(rt_coef%rdu(lay,i,1)*exu(1)*rt_coef%grud(lay-1,1,j) &
                   + rt_coef%rdu(lay,i,2)*exu(2)*rt_coef%grud(lay-1,2,j)) * exu(j)
          b(i,j) = (rt_coef%tdd(lay,i,1)*exu(1)*rt_coef%grud(lay-1,1,j) &
                  + rt_coef%tdd(lay,i,2)*exu(2)*rt_coef%grud(lay-1,2,j)) * exu(j)
        end do
        a(j,j) = (1.0d0, 0.0d0) + a(j,j)
      end do
      
      call mtxinv_2(a)
      
      do j = 1, 2
        do i = 1, 2
          rt_coef%gtu(lay,i,j) = a(i,1)*rt_coef%tuu(lay,1,j) + a(i,2)*rt_coef%tuu(lay,2,j)
        end do
      end do
      
      do j = 1, 2
        do i = 1, 2
          rt_coef%grud(lay,i,j) = rt_coef%rud(lay,i,j) + b(i,1)*rt_coef%gtu(lay,1,j) &
                              + b(i,2)*rt_coef%gtu(lay,2,j)
        end do
      end do
    end do
    
    ! (nly-1)界面的GRDU和GTDD
    lay = nly - 1
    hd = media%z(lay+1) - media%z(lay)
    exd(1) = exp(-cpn(lay+1)*hd)
    exd(2) = exp(-csn(lay+1)*hd)
    rt_coef%ex(lay,1) = exd(1)
    rt_coef%ex(lay,2) = exd(2)
    
    ! SH波
    rt_coef%grdu0(lay) = rt_coef%rdu0(lay)
    rt_coef%gtd0(lay) = rt_coef%tdd0(lay)
    
    ! PSV波
    do j = 1, 2
      do i = 1, 2
        rt_coef%gtd(lay,i,j) = rt_coef%tdd(lay,i,j)
        rt_coef%grdu(lay,i,j) = rt_coef%rdu(lay,i,j)
      end do
    end do
    
    ! lay=nly-2,...,ls界面的GRDU和GTDD
    do lay = nly-2, ls, -1
      hd = media%z(lay+1) - media%z(lay)
      exd(1) = exp(-cpn(lay+1)*hd)
      exd(2) = exp(-csn(lay+1)*hd)
      rt_coef%ex(lay,1) = exd(1)
      rt_coef%ex(lay,2) = exd(2)
      
      ! SH波
      tem1 = rt_coef%rud0(lay) * exd(2)
      tem2 = rt_coef%grdu0(lay+1) * exd(2)
      tem3 = rt_coef%tuu0(lay) * exd(2)
      rt_coef%gtd0(lay) = rt_coef%tdd0(lay) / (1.0d0 - tem1*tem2)
      rt_coef%grdu0(lay) = rt_coef%rdu0(lay) + tem3*tem2*rt_coef%gtd0(lay)
      
      ! PSV波
      do j = 1, 2
        do i = 1, 2
          a(i,j) = -(rt_coef%rud(lay,i,1)*exd(1)*rt_coef%grdu(lay+1,1,j) &
                   + rt_coef%rud(lay,i,2)*exd(2)*rt_coef%grdu(lay+1,2,j)) * exd(j)
          b(i,j) = (rt_coef%tuu(lay,i,1)*exd(1)*rt_coef%grdu(lay+1,1,j) &
                  + rt_coef%tuu(lay,i,2)*exd(2)*rt_coef%grdu(lay+1,2,j)) * exd(j)
        end do
        a(j,j) = (1.0d0, 0.0d0) + a(j,j)
      end do
      
      call mtxinv_2(a)
      
      do j = 1, 2
        do i = 1, 2
          rt_coef%gtd(lay,i,j) = a(i,1)*rt_coef%tdd(lay,1,j) + a(i,2)*rt_coef%tdd(lay,2,j)
        end do
      end do
      
      do i = 1, 2
        do j = 1, 2
          rt_coef%grdu(lay,i,j) = rt_coef%rdu(lay,i,j) + b(i,1)*rt_coef%gtd(lay,1,j) &
                              + b(i,2)*rt_coef%gtd(lay,2,j)
        end do
      end do
    end do
    
    rt_coef%ex(ls,1) = exp(-cpn(ls)*(media%z(ls)-media%z(ls-1)))
    rt_coef%ex(ls,2) = exp(-csn(ls)*(media%z(ls)-media%z(ls-1)))
  end subroutine grt_coefs
  
end module grt_coefficients